"""Ballistic forces and celestial body data"""

# %%
import math
from datetime import datetime

import numpy as np
from astropy import units as u
from astropy.constants import G, M_earth, R_earth
from astropy.time import Time
from orbital import elements_from_state_vector
from poliastro.earth.atmosphere import COESA76
from scipy.integrate import solve_ivp
from sgp4.api import Satrec

from mission_parameters import (
    C_F,
    JD_FR_START,
    JD_START,
    MIDDLE_AREA,
    SAT_DRY_MASS,
    TLE_LINE_1,
    TLE_LINE_2,
)
from sun import SUN_ANGLE

# G - Newtonian constant of gravitation 	6.67384e-11 m^3 kg^-1 s^-2
# M_earth 5.97216787e+24 kg Earth mass
# Гравитационная постоянная Земли
MU = G.value * M_earth.value


SIM_START_TIME = Time(JD_START, JD_FR_START, format="jd").value
# this calculation is CPU consuming, thus make it once and global

coesa76 = COESA76()
# We build the atmospheric instances


# %%
def sun_vector(current_time):
    """Sun unit vector in ECI

    Parameters
    ----------
    current_time : float
        current simulation time started from zero

    Returns
    -------
    r_earth_sun : 1x3 array
        sun position in ECI
    """
    # page 420 in book
    # Fundamentals of spacecraft attitude determination and control

    # CPU consuming with astropy library usage:
    # sim_start_time = Time(JD_START, JD_FR_START, format="jd")
    # JD = sim_start_time + current_time * u.second
    # JD = JD.value

    # Fast method, not using astropy library:
    JD = SIM_START_TIME + current_time / 86400

    t_ut1 = (JD - 2451545) / 36525

    phi_sun = math.radians(280.460 + 36000.771 * t_ut1)  # mean longitude
    phi_sun = phi_sun % (2 * math.pi)

    M_sun = math.radians(357.5277233 + 35999.05034 * t_ut1)  # mean anomaly
    M_sun = M_sun % (2 * math.pi)

    phi_ecliptic = (
        phi_sun
        + math.radians(1.914666471) * math.sin(M_sun)
        + 0.019994643 * math.sin(2 * M_sun)
    )  # longitude of the ecliptic
    eps_ecliptic = math.radians(
        23.439291 - 0.0130042 * t_ut1
    )  # obliquity of the ecliptic

    # r_earth_sun = get_sun(t).cartesian.xyz.to(u.m).value
    e_earth_sun = [
        math.cos(phi_ecliptic),
        math.cos(eps_ecliptic) * math.sin(phi_ecliptic),
        math.sin(eps_ecliptic) * math.sin(phi_ecliptic),
    ]

    return e_earth_sun


# %%
def sun_visibility(current_time, r_earth_sat):
    """Defines Sun visibility from satellite position.

    Parameters
    ----------
    current_time : float
        simulation time starting from zero, second.
    r_earth_sat - 1x3 numpy array
        vector of satellite coordinates in ECI, m.

    Returns
    -------
    float
        1. is for fully visible
        0. is for full eclipse
    """
    # https://en.wikipedia.org/wiki/Angular_diameter

    # r_earth_sun = sun_vector(current_time)
    # Position vector from the spacecraft to the Sun, GCI frame
    # r_sat_sun = r_earth_sun - r_earth_sat
    # e_sat_sun = r_sat_sun / np.linalg.norm(r_sat_sun)
    e_sat_sun = sun_vector(current_time)

    # calculations are made in satellite centered frame with axis parallel to ECI
    # TODO convert calculations to orbital frame, should be the same result
    r_earth_sat = np.array(r_earth_sat)
    r_sat_earth = -r_earth_sat
    e_sat_earth = r_sat_earth / np.linalg.norm(r_sat_earth)

    earth_angle = math.asin(R_earth.to(u.m).value / np.linalg.norm(r_earth_sat))
    s = math.acos(np.dot(e_sat_earth, e_sat_sun))

    # Calculate the size of the lune (http://mathworld.wolfram.com/Lune.html) in rad
    # Consider calculation in a plane on distance 1 meter from satellite
    # so we can can compare angels like radiuses
    if s < (earth_angle + SUN_ANGLE) and s > (earth_angle - SUN_ANGLE):
        lunedelta = 0.25 * math.sqrt(
            (SUN_ANGLE + earth_angle + s)
            * (earth_angle + s - SUN_ANGLE)
            * (s + SUN_ANGLE - earth_angle)
            * (SUN_ANGLE + earth_angle - s)
        )
        lune_area = (
            2 * lunedelta
            + SUN_ANGLE
            * SUN_ANGLE
            * (
                math.acos(
                    ((earth_angle * earth_angle) - (SUN_ANGLE * SUN_ANGLE) - (s * s))
                    / (2 * SUN_ANGLE * s)
                )
            )
            - earth_angle
            * earth_angle
            * (
                math.acos(
                    ((earth_angle * earth_angle) + (s * s) - (SUN_ANGLE * SUN_ANGLE))
                    / (2 * earth_angle * s)
                )
            )
        )
        res = lune_area / (math.pi * SUN_ANGLE * SUN_ANGLE)

    elif s <= (earth_angle - SUN_ANGLE):
        res = 0
    else:  # If s>earth_angle+SUN_ANGLE there is no eclipse taking place
        res = 1

    return res


# %%
def satellite_initial_orbit_parameters():
    """Calculates satellite initial position and speed from TLE

    Returns
    -------
    res
        [0:3] - position vector, m.
        [3:6] - velocity vector, m.
    """

    spacecraft = Satrec.twoline2rv(TLE_LINE_1, TLE_LINE_2)

    _, r_start, v_start = spacecraft.sgp4(JD_START, JD_FR_START)

    r_start = [element * 1e3 for element in r_start]  # convert from km to m
    v_start = [element * 1e3 for element in v_start]  # convert from km/s to m/s

    res = list(np.array([r_start, v_start]).flat)

    return res


# %%
def thrusting_zone_threshold(a, e):
    """
    Calculates allowed max angular difference between satellite true anomaly and
    periapsis (or apoapsis - depends on strategy) where thrusting is allowed.

    To circularize elliptical orbit thrusting allowed zone should be narrowed
    while delta_r = ra - rp growing.

    If delta_r is 0, thrusting should be allowed anywhere
    i.e. allowed true anomaly is [0,2pi).
    While delta_r -> inf, allowed angular zone around periapsis or apoapsis ->0
    it means thrusting is allowed only in one point.

    Parameters
    ----------
    a : float
        semimajor axis
    e : float
        eccentricity

    Returns
    -------
    float
    """
    # delta_r = ra-rp = a(1+e) - a(1-e)
    delta_r = 2 * a * e

    # coefficient are empirical
    res = math.pi * 0.6 ** (delta_r / 10e3)
    return res


# %%
def thrusting_allowed_event(t, y, valve_last_time_opened, sat_state, strategy):
    """In order to circularize elliptical orbit only part of the orbit allowed
    to apply thrusting.

    Parameters
    ----------
    y : [type]
        [description]
    strategy : string
        One of "increase periapsis" or "decrease apoapsis"

    Returns
    -------
    float
        >0 allowed, <0 not allowed, =0 - event point
    """

    r = np.array(y[2:5])
    v = np.array(y[5:8])
    orb_elements = elements_from_state_vector(r, v, MU)
    a, e, _, _, _, f = orb_elements
    # a - semimajor axis
    # e - eccentricity
    # f - true anomaly
    delta_f_max = thrusting_zone_threshold(a, e)
    if strategy == "increase periapsis":
        res = delta_f_max - abs(f - math.pi)
        # change state in 2 points around apoapsis: pi+/-delta_f_max
    else:
        res = abs(f - math.pi) - math.pi + delta_f_max
        # change state in 2 points around periapsis:
        # delta_f_max, 2pi-delta_f_max

    return res


# %%
def sat_is_in_thrusting_allowed_zone(y, strategy):
    """In order to circularize elliptical orbit only part of the orbit allowed
    to apply thrusting.

    Parameters
    ----------
    y : [type]
        [description]
    strategy : [type]
        [description]

    Returns
    -------
    boolean
    """
    # r = np.array(y[2:5])
    # v = np.array(y[5:8])
    # orb_elements = elements_from_state_vector(r, v, MU)
    # a, e, _, _, _, f = orb_elements
    # # a - semimajor axis
    # # e - eccentricity
    # # f - true anomaly

    # delta_f_max = thrusting_zone_threshold(a, e)

    # if strategy == "increase periapsis":
    #     res = math.pi - delta_f_max < f < math.pi + delta_f_max
    # else:
    #     res = (f < delta_f_max) or (f > 2 * math.pi - delta_f_max)
    res = thrusting_allowed_event(None, y, None, None, strategy=strategy) > 0

    return res


# %%
def gravity_acceleration(r_vector: np.ndarray) -> np.ndarray:
    """Calculates satellite's gravity acceleration vector.
    Reference: "Основы теории движения ИСЗ. Часть вторая: возмущенное движение", В.И.
    Крылов, page 11.

    Parameters
    ----------
    r_vector - 1x3 array of float
        vector of coordinates, m.
    Returns
    -------
    gravity_acc_vector - 1x3 array of float
        gravity acceleration vector, m/s^2.
    """

    r_vector = np.array(r_vector)

    x = r_vector[0]
    y = r_vector[1]
    z = r_vector[2]
    r = np.linalg.norm(r_vector)  # длина r_vector

    J_2 = 1082.636023 * 10**-6
    J_3 = -2.532435 * 10**-6
    J_4 = -1.6204 * 10**-6
    J_5 = -0.227716 * 10**-6
    J_6 = 0.00608 * 10**-6
    a_e = 6378136  # средний радиус Земли

    dr_J_2_vec = (
        3
        / 2
        * J_2
        * MU
        * a_e**2
        / (r**5)
        * np.array(
            [
                5 * x * z**2 / (r**2) - x,
                5 * y * z**2 / (r**2) - y,
                5 * z**3 / (r**2) - 3 * z,
            ]
        )
    )
    dr_J_3_vec = (
        5
        / 2
        * J_3
        * MU
        * a_e**3
        / (r**7)
        * np.array(
            [
                7 * x * z**3 / (r**2) - 3 * x * z,
                7 * y * z**3 / (r**2) - 3 * y * z,
                7 * z**4 / (r**2) - 3 * z**2 + 3 / 5 * r**2,
            ]
        )
    )
    dr_J_4_vec = (
        1
        / 4
        * J_4
        * MU
        * a_e**4
        / (r**7)
        * np.array(
            [
                315 * x * z**4 / (2 * r**4)
                - 105 * x * z**2 / (r**2)
                + 15 * x / 2,
                315 * y * z**4 / (2 * r**4)
                - 105 * y * z**2 / (r**2)
                + 15 * y / 2,
                315 * z**5 / (2 * r**4) - 175 * z**3 / (r**2) + 75 * z / 2,
            ]
        )
    )
    dr_J_5_vec = (
        1
        / 8
        * J_5
        * MU
        * a_e**5
        / (r**9)
        * np.array(
            [
                693 * x * z**5 / (2 * r**4)
                - 630 * x * z**3 / (r**2)
                + 105 * x * z,
                693 * y * z**5 / (2 * r**4)
                - 630 * y * z**3 / (r**2)
                + 105 * y * z,
                (693 * z**6 - 945 * z**3 * r**2) / (r**4)
                + 315 * z**2
                - 15 * r**2,
            ]
        )
    )
    dr_J_6_vec = (
        1
        / 16
        * J_6
        * MU
        * a_e**6
        / (r**9)
        * np.array(
            [
                3003 * x * z**6 / (r**6)
                - 3465 * x * z**4 / (r**4)
                + 945 * x * z**2 / (r**2)
                - 35 * x,
                3003 * y * z**6 / (r**6)
                - 3465 * y * z**4 / (r**4)
                + 945 * y * z**2 / (r**2)
                - 35 * y,
                3003 * z**7 / (r**6)
                - 4851 * z**5 / (r**4)
                + 2205 * z**3 / (r**2)
                - 245 * z,
            ]
        )
    )

    gravity_acc_vector = (
        -G * M_earth * r_vector / np.linalg.norm(r_vector) ** 3
    ).value + (dr_J_2_vec + dr_J_3_vec + dr_J_4_vec + dr_J_5_vec + dr_J_6_vec)

    return gravity_acc_vector


# %%
def satellite_acceleration(
    r_vector: np.ndarray,
    thrust_v_i: np.ndarray,
    m_fuel: float,
    v_vector: np.ndarray,
) -> np.ndarray:
    """Calculates satellite's acceleration

    Parameters
    ----------
    r_vector - 1x3 array of float
        vector of coordinates, m.
    thrust_v_i - 1x3 array of float
        vector of thrust in inertial frame, N.
    m_fuel - float
        fuel mass, kg.
    v_vector - 1x3 array of float
        vector of velocity, m/s

    Returns
    -------
    a_vector - 1x3 array of float
        acceleration vector, m/s^2.
    """

    r_vector = np.array(r_vector)
    thrust_v_i = np.array(thrust_v_i)
    v_vector = np.array(v_vector)

    sat_mass = SAT_DRY_MASS + m_fuel

    height = (np.linalg.norm(r_vector) - R_earth.value) * u.m

    # FIXME работает только в небольшом диапазоне высот от 0 до 1000 км. Поддержка очень
    # большого расстояния сделана для того, чтобы проходили unit-тесты для бесконечного
    # удаленного КА.
    #
    # Анализ данных прибора SWAN на космическом аппарате SOHO показал, что самая внешняя
    # часть экзосферы Земли (геокорона) простирается примерно на 100 радиусов Земли или
    # около 640 тыс. км, то есть гораздо дальше орбиты Луны
    # https://ru.wikipedia.org/wiki/%D0%90%D1%82%D0%BC%D0%BE%D1%81%D1%84%D0%B5%D1%80%D0%B0_%D0%97%D0%B5%D0%BC%D0%BB%D0%B8
    if np.linalg.norm(height.value) >= 640_000_000:
        drag_force = np.array([0, 0, 0])
    # coesa76 считает плотность атмосферы только до 1000 км
    else:
        drag_force = (
            -C_F
            * MIDDLE_AREA
            * v_vector
            * np.linalg.norm(v_vector)
            * coesa76.density(height).value
            / 2
        )

    sum_forces = gravity_acceleration(r_vector) * sat_mass + thrust_v_i + drag_force

    acc_vector = sum_forces / sat_mass

    return acc_vector


# %%
def f_ballistic_t_y(t, y):
    """Right side of ODE
    dy/dt = f_t_y.
    Position and velocity derivatives

    Parameters
    ----------
    t : float
        time, sec.
    y : array like
        [r1x,r1y,r1z,
        v1x,v1y,v1z,
        ...
        rix,riy,riz,
        vix,viy,viz]
        coordinate vector components, m;
        velocity vector components, m/sec.

    Returns
    -------
    res : array of float
        res[0:3] - velocity vector components for 1st satellite, m/s
        res[3:6] - acceleration vector components for 1st satellite, m/s^2
        res[6:12] - same for 2nd satellite
        ...
        res[6*(i-1):6*i] - same for satellite #i
    """

    res = []
    for i in range(round(len(y) / 6)):
        ri_vector = np.array(y[6 * i : 6 * i + 3])
        vi_vector = y[6 * i + 3 : 6 * i + 6]
        ai_vector = -G * M_earth * ri_vector / np.linalg.norm(ri_vector) ** 3
        ai_vector = ai_vector.value
        res = np.concatenate([res, vi_vector, ai_vector])

    return res


# %%
def phasing_angle(r1, r2):
    """Phasing event reached when satellites are 1/4 of one rotation or
    the dot product is equal to zero. Multiply this time by number 4 will be
    enough for 2 satellites have roughly one rotation difference. It means
    that any other satellite in this time can be placed at any phasing angle
    between this two satellites.

    Parameters
    ----------
    t : float
        time, sec
    r1 : [array of float] coordinate vector components, m;
    r2 : [array of float] coordinate vector components, m;

    Return
    ------
    Dot product of r1 and r2 unit vectors
    """

    r1_unit = r1 / np.linalg.norm(r1)
    r2_unit = r2 / np.linalg.norm(r2)

    return np.dot(r1_unit, r2_unit)


# %%
def free_flight_simulation(y0, t0, t_end):
    """Non thrusting flight of 2 satellites simulation
    until their rendezvous,
    probably on different altitudes but on the same (almost) zenith vector.

    Parameters
    ----------
    y0 : array of float
        [r1x,r1y,r1z,
        v1x,v1y,v1z,
        ...
        rix,riy,riz,
        vix,viy,viz,]
        coordinate vector components, m;
        velocity vector components, m/sec.
    t0 : float
        simulation start time, sec
    t_end : float
        simulation end time, sec

    Returns
    -------
    Bunch object with the following fields defined:
        t : ndarray, shape (n_points,) Time points.
        y : ndarray, shape (n, n_points) Values of the solution at t
    """
    t_span = [t0, t_end]
    sol_next_steps = solve_ivp(
        f_ballistic_t_y,
        t_span,
        y0,
        method="Radau",
        rtol=1e-4,
    )
    return sol_next_steps


# %%
def free_flight_phasing(t0, sat_y_last_points, alpha_goal):
    r1 = sat_y_last_points[0][2:5]
    v1 = sat_y_last_points[0][5:8]
    r2 = sat_y_last_points[1][2:5]
    v2 = sat_y_last_points[1][5:8]

    # https://space.stackexchange.com/questions/1904/how-to-programmatically-calculate-orbital-elements-using-position-velocity-vecto
    # https://github.com/RazerM/orbital/blob/0.7.0/orbital/utilities.py#L252
    # this follows from the method laid out in Fundamentals of Astrodynamics and
    # Applications, by Vallado, 2007
    orb_elements1 = elements_from_state_vector(r1, v1, MU)
    orb_elements2 = elements_from_state_vector(r2, v2, MU)
    a1 = orb_elements1[0]
    a2 = orb_elements2[0]
    # https://en.wikipedia.org/wiki/Orbital_period
    T1 = 2 * math.pi * math.sqrt(a1**3 / MU)
    T2 = 2 * math.pi * math.sqrt(a2**3 / MU)
    print(f"{T1=}\n{T2=}")

    y0 = np.concatenate((r1, v1, r2, v2), axis=0)

    # Initial phasing difference:
    alpha_0 = math.acos(phasing_angle(r1, r2))
    # Obtained from equation:
    # T_min(n+(alpha_goal-2*alpha0)/(2pi)) = T_max*n
    t_end = t0 + T1 * T2 / abs(T1 - T2) * (alpha_goal - 2 * alpha_0) / (2 * math.pi)

    sim_start_time = datetime.now()
    print(f"\n\nFree flight phasing started at {sim_start_time}")
    free_flight_phasing_sol = free_flight_simulation(y0, t0, t_end)
    sim_end_time = datetime.now()
    print(f"Free flight phasing finished at {sim_end_time}")

    # freeze all parameters except r1, v2, r2, v2
    r1 = free_flight_phasing_sol.y[0:3][:, -1]
    v1 = free_flight_phasing_sol.y[3:6][:, -1]
    r2 = free_flight_phasing_sol.y[6:9][:, -1]
    v2 = free_flight_phasing_sol.y[9:12][:, -1]
    sat_y_last_points[0][2:5] = r1
    sat_y_last_points[0][5:8] = v1
    sat_y_last_points[1][2:5] = r2
    sat_y_last_points[1][5:8] = v2

    return free_flight_phasing_sol, sat_y_last_points
